#import "MyProcessWorker.h"

#include <mach/mach_types.h>
#include <mach/mach_traps.h>
#include <mach/mach_error.h>

#include <sys/sysctl.h>
#include <sys/reboot.h>
#include <unistd.h>

@interface MyProcessWorker(PrivateMethods)

- (void)notifyCrash;

@end

@implementation MyProcessWorker

- (id)init
{
	if (self = [super init])
	{
		[self walkProcessList];
	}
	
	return self;
}

- (void)walkProcessList
{
	int					mib[4];
	int 				i;
	size_t 				bufSize;
	int 				nentries;
	struct kinfo_proc 	*kprocbuf = NULL, *kp;
	BOOL				initialWalk = (myLastSeenPID == 0);
	
	NS_DURING
	{
		mib[0] = CTL_KERN;
		mib[1] = KERN_PROC;
		mib[2] = KERN_PROC_ALL;
		mib[3] = 0;
	
		if (sysctl(mib, 4, NULL, &bufSize, NULL, 0) < 0)
		{
			[NSException raise:NSInternalInconsistencyException
						format:@"unable to count processes (sysctl)"];
		}
		
		nentries = bufSize / sizeof(struct kinfo_proc);
		kprocbuf = kp = (struct kinfo_proc *)NSZoneMalloc(NULL, bufSize);

		if (sysctl(mib, 4, kp, &bufSize, NULL, 0) < 0)
		{
			[NSException raise:NSInternalInconsistencyException
						format:@"unable to get process list (sysctl)"];
		}   
		
		// it lies the first time, fortunately with a buffer that is too big
		nentries = bufSize / sizeof(struct kinfo_proc);
	
		if (kp != NULL)
		{
			for(i=nentries; --i>=0; ++kp)
			{
				if (kp->kp_proc.p_pid > myLastSeenPID)
				{
					myLastSeenPID = kp->kp_proc.p_pid;

					if ( ! initialWalk)
					{
						[NSThread detachNewThreadSelector:@selector(myThreadMain:)
							toTarget:self withObject:[NSNumber numberWithInt:myLastSeenPID]];
					}
				}
			}
		}
	}
	NS_HANDLER
	{
		NSLog(@"Couldn't install exception handler for pid: %d; Exception :%@", [localException description]);
	}
	NS_ENDHANDLER;
	
	if (kprocbuf != NULL)
		NSZoneFree(NULL, kprocbuf);
}

#define MACH_CHECK_ERROR(name,ret)	\
		if (ret != KERN_SUCCESS) {	\
		mach_error(#name, ret);		\
		exit(1);					\
} 

- (void)myThreadMain:(id)arg
{
	const int 			MSG_SIZE = 512;
	id		  		 	 pid = [arg retain];
    mach_msg_header_t 	*msg, *reply;
    kern_return_t 		krc;
	mach_port_t 		_exceptionPort;
	task_t 				targetTask;
	
	NS_DURING
	{
		krc = mach_port_allocate(mach_task_self(), MACH_PORT_RIGHT_RECEIVE, &_exceptionPort);
		if (krc != KERN_SUCCESS)
		{
			[NSException raise: NSInternalInconsistencyException format: @"Unable to create handler port, krc = %d, %s", krc, mach_error_string(krc)];
		}
		
		krc = mach_port_insert_right(mach_task_self(), _exceptionPort, _exceptionPort, MACH_MSG_TYPE_MAKE_SEND);
		if (krc != KERN_SUCCESS)
		{
			[NSException raise: NSInternalInconsistencyException format: @"Unable insert send write into handler port, krc = %d, %s", krc, mach_error_string(krc)];
		}
	
		krc = task_for_pid(mach_task_self(), [pid intValue], &targetTask);
		if (krc != KERN_SUCCESS)
		{
			[NSException raise: NSInternalInconsistencyException format: @"task_for_pid failed, krc = %d, %s", krc, mach_error_string(krc)];
		}
		
		krc = task_set_exception_ports(targetTask, EXC_MASK_BAD_ACCESS | EXC_MASK_BAD_INSTRUCTION | EXC_MASK_ARITHMETIC, _exceptionPort, EXCEPTION_DEFAULT, THREAD_STATE_NONE);
		if (krc != KERN_SUCCESS)
		{
			[NSException raise: NSInternalInconsistencyException format: @"task_set_exception_ports failed, krc = %d, %s", krc, mach_error_string(krc)];
		}
	
		msg = alloca(MSG_SIZE);
		reply = alloca(MSG_SIZE);

		while (TRUE)
		{
			NSAutoreleasePool	*pool = [[NSAutoreleasePool alloc] init];
			
			// we only get what we consider to be fatal messages,
			// so just shoot the process and do our thing

			krc = mach_msg(msg, MACH_RCV_MSG, MSG_SIZE, MSG_SIZE, _exceptionPort, 0, MACH_PORT_NULL);
			if (krc == KERN_SUCCESS)
			{
				kill([pid intValue], SIGKILL);
				[self notifyCrash];
			}

			[pool release];
		}
	}
	NS_HANDLER
	{
	
	}
	NS_ENDHANDLER;
}

@end

@implementation MyProcessWorker(PrivateMethods)

- (void)notifyCrash
{
	const int 		FREE_CRASH_LIMIT = 1;
	CFStringRef 	APPLICATION_PREF_ID = CFSTR("Unprotected Memory");
	CFStringRef 	ALREADY_REBOOTED_KEY = CFSTR("AlreadyRebooted");
	
	BOOL			noGrace = NO;
	CFTimeInterval	timeout = 5;

	myCrashCount++;
	
	if ( CFPreferencesGetAppBooleanValue(ALREADY_REBOOTED_KEY, APPLICATION_PREF_ID, NULL) )
	{
		noGrace = YES;
		timeout = 0.5;
	}
	
	if (noGrace || (myCrashCount > FREE_CRASH_LIMIT))
	{
		SInt32			response;
		CFOptionFlags	responseFlags;
	
		response = CFUserNotificationDisplayAlert(timeout, 0L, NULL, NULL, NULL,
													CFSTR("An application has unexpectedly quit."),
													CFSTR("You have 5 seconds to save all your work before the machine reboots."),
													CFSTR("Restart"),
													NULL,
													NULL,
													&responseFlags);
													
		CFPreferencesSetAppValue(ALREADY_REBOOTED_KEY, kCFBooleanTrue, APPLICATION_PREF_ID);
		CFPreferencesAppSynchronize(APPLICATION_PREF_ID);

													
		if (0 != reboot(RB_AUTOBOOT))
		{
			// call only returns if it could fail, otherwise
			// the machine is rebooting
			
			NSLog(@"reboot didn't work, did you run it as root dummy?");
		}
	}
	else
	{
		(void)CFUserNotificationDisplayNotice(0, 0L, NULL, NULL, NULL,
											  CFSTR("An application has unexpectedly quit."),
											  CFSTR("Aren't you glad you didn't have to reboot your machine?"),
											  CFSTR("Whew!"));
	}
}

@end
